#!/usr/bin/python
"""Generates samples based on given two criterias; based on threshold, and based on sampling method.

Based on threshold. Energy matrix (m) and threshold (t) is required along with filename.

Based on method. Energy matrix (m), Method name(s) and number of samples (n) is required along with filename.
"""

import sys
import getopt
import types

from pdf import fermi_dirac
from pdf import boltzmann

# -- globals --
cutoff = None      # cutoff threshold
matrix = None      # energy matrix
method = None      # sampling method, options are fermi-dirac (f) and boltzmann (b)
n_samples = None   # number of samples required
debug = False      # debug off
best = None        # Best sequence having highest affinity
seq_len = None     # Length of each sequence
positive_seqs = [] # All sequences doing better than or equal to threshold
base_map = {'a':0, 'c':1, 'g':2, 't':3, 'A':0, 'C':1, 'G':2, 'T':3}

def usage():
    print '''Usage : sample.py [OPTIONS]... [FILE]

    [FILE] is file containing energy matrix
    
    Options:
    --------
    -h, --help                   Display this text and exit.
    -d, --debug                  Debug 0n
    -t, --threshold[=real]       Energy cutoff threshols.
    -s, --sampling[=f|b]         Type of sampling; either fermi-dirac or boltzmann.
    -n, --samples[=integer]      Total number of samples required.
     '''

def main(argv, energy_matrix=None):
    global cutoff
    global method
    global n_samples
    global matrix
    global debug
    global base_map
    global seq_len
    global positive_seqs
    
    if len(argv) == 0: usage(); sys.exit(2)

    try:
        opts, args = getopt.getopt(argv, "hdt:s:n:", ["help", "debug", "threshold", "sampling", "samples"])
    except getopt.GetoptError:
        usage()
        sys.exit(2)

    if len(args) == 0:
        if energy_matrix == None:
            sys.exit("Please specify matrix file.")
        else:
            matrix = energy_matrix
    
    for opt, arg in opts:
        if opt in ("-h", "--help"):
            usage()
            sys.exit(2)
        elif opt in ("-t", "--threshold"):
            try:
                cutoff = float(arg)
            except ValueError:
                print "Invalid value for '-t'"         
        elif opt in ("-d", "--debug"):
            debug = True
        elif opt in ("-s", "--sampling"):
            method = arg
            if method not in ("f", "b"):
                usage()
                sys.exit(4)
        elif opt in ("-n", "--samples"):
            try:
                n_samples = int(arg)
            except ValueError:
                print "Invalid value for '-n'"
        else:
            usage()
            sys.exit()

    if energy_matrix == None:
        matrix = file2mat(args[0])
    else:
        matrix = energy_matrix
        
    seq_len = len(matrix[0])

    #if (debug): print "cutoff=%f, method=%s, total=%d" % (cutoff, method, n_samples)

    if debug: showMatrix(matrix)

    expand('a')
    expand('c')
    expand('g')
    expand('t')
        
    if debug: print "Total positive seqs", len(positive_seqs)
    # return after printing all sequences having energy better than threshold
    # only if there is no method (nothing to be done for sampling, neither FD nor BM)
    if not method:
    #    for s in positive_seqs:
    #        print s
        return positive_seqs

    if n_samples == None:
        sys.exit("Please specify number of samples required.")
    
    if method == "f": # fermi-dirac sampling
        samples = fermi_dirac(positive_seqs, n_samples)
    else:             # boltzmann sampling
        samples = boltzmann(matrix, positive_seqs, n_samples)

    #for s in samples:
    #    print s

    return samples

def file2mat(filename):
    lines = file2lines(filename)
    matrix = [[1 for x in xrange(len(lines[0].split()))] for y in xrange(4)]

    for i in xrange(4):
        line = (lines[i]).split()
        for j in xrange(len(line)):
            matrix[i][j] = float(line[j])

    return matrix

def file2lines(filename):
    # open file having alignments
    try:
        file = open(filename, 'r')
    except IOError:
        print "No such file '%s'" % filename
        sys.exit(2)
    except:
        print "Unexpected error:", sys.exc_info()[0]
        sys.exit(2)

    lines = file.readlines()

    for i in xrange(len(lines)):
        lines[i] = (lines[i]).strip()

    # return a list cotaining all the lines in given file
    return lines

def showMatrix(weight_mat):
    for row in weight_mat:
        for col in row:
            if (type(col) == types.IntType):
                print col,
            elif (type(col) == types.FloatType):
                print '%.2f' % col,
            else:
                print "Matrix element type not supported."
                sys.exit(2)
        print

def score(line):
    global matrix
    global base_map
    
    score = 0

    for i in xrange(len(line)):
        score += matrix[base_map[line[i]]][i]

    return score

def best_seq():
    global best
    
    if best != None:
        return best
    
    global debug
    
    max_idx = [0 for i in xrange(len(matrix[0]))]

    for i in xrange(len(max_idx)):
        max_idx[i] = 0
        for j in xrange(4):
            if matrix[j][i] > matrix[max_idx[i]][i]:
                max_idx[i] = j
    if debug:
        for i in max_idx:
            print i,
        print

    seq = ''
    for i in max_idx:
        if i == 0:
            seq += 'a'
        elif i ==1:
            seq += 'c'
        elif i == 2:
            seq += 'g'
        elif i == 3:
            seq += 't'
        else:
            sys.exit("Unknown base in sequence. Index out of 0-3")

    best = seq
    return seq

def expand(seq):
    global seq_len
    global debug
    global positive_seqs

    best = best_seq()
    
    if len(seq) == seq_len:
        if cutoff==None:
            positive_seqs.append(seq)
        else:
            if score(seq) >= cutoff:
                positive_seqs.append(seq)
                #print score(seq)
            else:
                pass
    elif len(seq) < seq_len:
        overestimate = seq
        
        for i in xrange(len(seq), seq_len):
            overestimate += best[i]

        if score(overestimate) >= cutoff:
            expand(seq+'a')
            expand(seq+'c')
            expand(seq+'g')
            expand(seq+'t')
        else:
            pass
    else:
        sys.exit("Sequence length more than limit. Some bug in logic.")

if __name__ == "__main__":
    seqs = main(sys.argv[1:])

    for seq in seqs:
        print seq
